package com.feonix.common.datasource;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

import javax.sql.DataSource;

import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.JndiDataSourceLookup;
import org.springframework.util.Assert;

public class AdaptiveDataSource extends AbstractDataSource implements InitializingBean {

	private Object defaultTargetDataSource;
	private boolean lenientFallback = true;

	private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();

	private Map<Object, DataSource> resolvedDataSources = new HashMap<Object, DataSource>();
	private DataSource resolvedDefaultDataSource;

	public AdaptiveDataSource() {
		DataSourceManager.setAdaptiveDataSource(this);
		DataSourceManager.loadDataSource();
	}

	public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
		this.defaultTargetDataSource = defaultTargetDataSource;
	}

	public void setLenientFallback(boolean lenientFallback) {
		this.lenientFallback = lenientFallback;
	}

	public void setDataSourceLookup(DataSourceLookup dataSourceLookup) {
		this.dataSourceLookup = new JndiDataSourceLookup();
	}

	public void afterPropertiesSet() {
	}

	public void add(Object key, Object object) {
		Object lookupKey = resolveSpecifiedLookupKey(key);
		DataSource dataSource = resolveSpecifiedDataSource(object);
		this.resolvedDataSources.put(lookupKey, dataSource);
	}

	public void setDefault(Object dataSourceLookup) {
		if (dataSourceLookup != null) {
			this.defaultTargetDataSource = dataSourceLookup;
			this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
		}
	}

	public void remove(Object key) {
		Object dataSourceKey = resolveSpecifiedLookupKey(key);
		if (this.resolvedDataSources.containsKey(dataSourceKey)){
			this.resolvedDataSources.remove(dataSourceKey);
		}
	}

	protected Object resolveSpecifiedLookupKey(Object dataSource) {
		return dataSource;
	}

	protected DataSource resolveSpecifiedDataSource(Object lookupKey) throws IllegalArgumentException {
		if (lookupKey instanceof DataSource)
			return ((DataSource) lookupKey);
		if (lookupKey instanceof String) {
			return this.dataSourceLookup.getDataSource((String) lookupKey);
		}
		throw new IllegalArgumentException(
				"Illegal data source value - only [javax.sql.DataSource] and String supported: " + lookupKey);
	}

	public Connection getConnection() throws SQLException {
		Connection connection = determineTargetDataSource().getConnection();
		connection.setAutoCommit(true);
		return connection;
	}

	public Connection getConnection(String username, String password) throws SQLException {
		Connection connection = determineTargetDataSource().getConnection(username, password);
		connection.setAutoCommit(true);
		return connection;
	}

	public <T> T unwrap(Class<T> iface) throws SQLException {
		if (iface.isInstance(this)) {
			return (T) this;
		}
		return determineTargetDataSource().unwrap(iface);
	}

	public boolean isWrapperFor(Class<?> iface) throws SQLException {
		return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
	}

	protected DataSource determineTargetDataSource() {
		Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
		Object lookupKey = determineCurrentLookupKey();
		DataSource dataSource = (DataSource) this.resolvedDataSources.get(lookupKey);
		if ((dataSource == null) && (this.lenientFallback || lookupKey == null)) {
			dataSource = this.resolvedDefaultDataSource;
		}
		if (dataSource == null) {
			throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
		}
		return dataSource;
	}

	protected Object determineCurrentLookupKey() {
		return DataSourceSwitch.getAlias();
	}

}
